#!/usr/bin/env python3
"""
完整Pipeline脚本

一键运行完整流程：
1. 数据预处理
2. 模型训练
3. 模型评估
4. 结果可视化
5. 报告生成
"""

import argparse
import sys
from pathlib import Path
import subprocess
from datetime import datetime

# 添加项目路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.utils.logger import get_logger


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='运行完整Pipeline')
    parser.add_argument('--config', type=str, default='configs/training.yaml',
                       help='配置文件路径')
    parser.add_argument('--skip_preprocessing', action='store_true',
                       help='跳过数据预处理步骤')
    parser.add_argument('--skip_training', action='store_true',
                       help='跳过训练步骤')
    parser.add_argument('--skip_evaluation', action='store_true',
                       help='跳过评估步骤')
    parser.add_argument('--checkpoint', type=str, default=None,
                       help='用于评估的检查点路径（如果跳过训练）')
    parser.add_argument('--output_dir', type=str, default='results',
                       help='输出目录')
    return parser.parse_args()


def run_preprocessing(config_path: str, logger):
    """运行数据预处理"""
    logger.info("=" * 60)
    logger.info("步骤1: 数据预处理")
    logger.info("=" * 60)
    
    try:
        # 运行预处理脚本
        cmd = [sys.executable, 'scripts/run_preprocessing.py', '--config', config_path]
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        logger.info("数据预处理完成")
        logger.info(result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"数据预处理失败: {e}")
        logger.error(e.stderr)
        return False


def run_training(config_path: str, output_dir: str, logger):
    """运行模型训练"""
    logger.info("=" * 60)
    logger.info("步骤2: 模型训练")
    logger.info("=" * 60)
    
    try:
        # 运行训练脚本，实时显示输出
        cmd = [
            sys.executable, 'scripts/train_unified_prompt.py',
            '--config', config_path,
            '--output_dir', str(Path(output_dir) / 'training')
        ]
        logger.info(f"执行命令: {' '.join(cmd)}")
        logger.info("开始训练，输出将实时显示...")
        
        # 使用Popen实时显示输出
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,  # 行缓冲
            universal_newlines=True
        )
        
        # 实时输出（逐行读取）
        import sys as sys_module
        for line in iter(process.stdout.readline, ''):
            if not line:
                break
            sys_module.stdout.write(line)
            sys_module.stdout.flush()
        
        # 等待进程完成
        process.wait()
        
        if process.returncode != 0:
            raise subprocess.CalledProcessError(process.returncode, cmd)
        
        logger.info("模型训练完成")
        
        # 查找最佳模型检查点
        checkpoint_dir = Path(output_dir) / 'training'
        best_checkpoint = checkpoint_dir / 'best_model.pt'
        
        if best_checkpoint.exists():
            logger.info(f"最佳模型检查点: {best_checkpoint}")
            return str(best_checkpoint)
        else:
            logger.warning("未找到最佳模型检查点")
            return None
    except subprocess.CalledProcessError as e:
        logger.error(f"模型训练失败: {e}")
        return None
    except Exception as e:
        logger.error(f"模型训练出错: {e}")
        return None


def run_evaluation(checkpoint_path: str, config_path: str, output_dir: str, logger):
    """运行模型评估"""
    logger.info("=" * 60)
    logger.info("步骤3: 模型评估")
    logger.info("=" * 60)
    
    if not checkpoint_path:
        logger.error("未提供检查点路径，无法进行评估")
        return False
    
    try:
        # 从配置中获取测试数据目录
        import yaml
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        test_data_dir = config.get('data', {}).get('test_dir', 'data/test')
        
        # 运行评估脚本
        cmd = [
            sys.executable, 'scripts/evaluate_unified_prompt.py',
            '--checkpoint', checkpoint_path,
            '--config', config_path,
            '--test_data_dir', test_data_dir,
            '--output_dir', str(Path(output_dir) / 'evaluation')
        ]
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        logger.info("模型评估完成")
        logger.info(result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"模型评估失败: {e}")
        logger.error(e.stderr)
        return False
    except Exception as e:
        logger.error(f"模型评估出错: {e}")
        return False


def main():
    """主函数"""
    args = parse_args()
    
    # 设置日志
    logger = get_logger("FullPipeline")
    logger.info("=" * 60)
    logger.info("开始运行完整Pipeline")
    logger.info("=" * 60)
    logger.info(f"配置文件: {args.config}")
    logger.info(f"输出目录: {args.output_dir}")
    
    start_time = datetime.now()
    
    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 步骤1: 数据预处理
    if not args.skip_preprocessing:
        success = run_preprocessing(args.config, logger)
        if not success:
            logger.error("数据预处理失败，终止Pipeline")
            return
    else:
        logger.info("跳过数据预处理步骤")
    
    # 步骤2: 模型训练
    checkpoint_path = args.checkpoint
    if not args.skip_training:
        checkpoint_path = run_training(args.config, args.output_dir, logger)
        if not checkpoint_path:
            logger.error("模型训练失败，终止Pipeline")
            return
    else:
        logger.info("跳过训练步骤")
        if not checkpoint_path:
            logger.error("跳过训练但未提供检查点路径，无法进行评估")
            return
    
    # 步骤3: 模型评估
    if not args.skip_evaluation:
        success = run_evaluation(checkpoint_path, args.config, args.output_dir, logger)
        if not success:
            logger.error("模型评估失败")
    else:
        logger.info("跳过评估步骤")
    
    # 完成
    end_time = datetime.now()
    duration = end_time - start_time
    
    logger.info("=" * 60)
    logger.info("Pipeline执行完成！")
    logger.info(f"总耗时: {duration}")
    logger.info(f"输出目录: {output_dir}")
    logger.info("=" * 60)


if __name__ == '__main__':
    main()
